import os
import pandas as pd
import matplotlib.pyplot as plt
from datetime import datetime, timedelta
import matplotlib as mpl

# 设置默认字体为支持中文的字体
mpl.rcParams['font.sans-serif'] = ['SimHei']  # 或者 'Microsoft YaHei', 'FangSong' 等
mpl.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

def calculate_equal_weighted_price(df):
    total_price = 0.0
    count = 0
    for index, row in df.iterrows():
        close_price = row['close']
        if pd.notna(close_price):  # Check if close_price is not NaN
            total_price += close_price
            count += 1
    return total_price / count if count > 0 else 0


def plot_price_curve(timestamps, equal_weighted_prices, btc_prices):
    # Ensure timestamps are valid datetime objects
    valid_timestamps = []
    for ts in timestamps:
        try:
            dt = pd.to_datetime(ts)
            valid_timestamps.append(dt)
        except ValueError:
            print(f"Invalid timestamp: {ts}")

    # Convert valid timestamps to matplotlib date format
    valid_timestamps = mpl.dates.date2num(valid_timestamps)

    fig, ax1 = plt.subplots(figsize=(20, 5))

    # Plot Equal Weighted Price on the left y-axis
    ax1.plot(valid_timestamps, equal_weighted_prices, linestyle='-', color='b', label='Equal Weighted Price')
    ax1.set_xlabel('Time')
    ax1.set_ylabel('Equal Weighted Price', color='b')
    ax1.tick_params(axis='y', labelcolor='b')
    ax1.grid(True)

    # Format x-axis to display dates in 'YYYY/MM/DD' format
    date_fmt = mpl.dates.DateFormatter('%Y/%m/%d')
    ax1.xaxis.set_major_formatter(date_fmt)
    plt.xticks(rotation=45)

    # Create a second y-axis for BTC Price
    ax2 = ax1.twinx()
    ax2.plot(valid_timestamps, btc_prices, linestyle='-', color='r', label='BTC Price')
    ax2.set_ylabel('BTC Price', color='r')
    ax2.tick_params(axis='y', labelcolor='r')

    # Add legends
    fig.legend(loc="upper left", bbox_to_anchor=(0, 1), bbox_transform=ax1.transAxes)

    plt.title('现货等权价格指数')
    plt.tight_layout()
    plt.show()

# Directory containing the CSV files
#data_dir = 'C:\\pythonproject\\coin\\data\\spot_binance_1h'

data_dir = 'D:\stock\data\coin-binance-candle-csv-1h'

# List to store dataframes
dfs = []
# Calculate the date 200 days ago from today
today = datetime.today()
date_200_days_ago = today - timedelta(days=365*3)

count = 0
# Read all CSV files in the directory
for filename in os.listdir(data_dir):
    if filename.endswith('.csv'):
        print(filename)

        file_path = os.path.join(data_dir, filename)
        df = pd.read_csv(file_path, parse_dates=['candle_begin_time'],skiprows=1,encoding='gbk')
        if 'symbol' not in df.columns:
            print(f"Warning: 'symbol' column not found in {filename}. Skipping this file.")
            continue
         # Filter data to keep only the last 200 days
        df = df[df['candle_begin_time'] >= date_200_days_ago]
        if df.size == 0:
            continue

        # 使用 resample 进行数据重采样，并指定聚合规则
        rule_dict = {'open': 'first', 'high': 'max', 'low': 'min', 'close': 'last', 'symbol': 'last'}
        resampled_df = df.resample(rule='1D', on = 'candle_begin_time').agg(rule_dict)

        # 使用 fillna 进行前向填充
        resampled_df = resampled_df.fillna(method='ffill')
        # 重置索引，将 'candle_begin_time' 转换回列
        resampled_df = resampled_df.reset_index().rename(columns={'index': 'candle_begin_time'})

        dfs.append(resampled_df)


# Combine all dataframes into one
print("数据合并中...")
combined_df = pd.concat(dfs)
print("数据合并分组中...")
# Calculate equal weighted price for each day using groupby
daily_combined_df = combined_df.groupby(combined_df['candle_begin_time'].dt.date)
equal_weighted_prices = daily_combined_df.apply(calculate_equal_weighted_price).tolist()

# Get the latest 200 days
timestamps = combined_df['candle_begin_time'].unique().tolist()
equal_weighted_prices = equal_weighted_prices

# Calculate BTC price for each day
btc_df = combined_df[combined_df['symbol'] == 'BTC-USDT']
btc_prices = btc_df['close'].tolist()[-200:]


plot_price_curve(timestamps, equal_weighted_prices, btc_prices)